from datetime import datetime
from typing import List, Optional
from uuid import UUID

from pydantic import BaseModel, Field, validator, root_validator


class MemoryBase(BaseModel):
    content: str
    metadata_: Optional[dict] = Field(default_factory=dict)

class MemoryCreate(MemoryBase):
    user_id: UUID
    app_id: UUID


class Category(BaseModel):
    name: str


class App(BaseModel):
    id: UUID
    name: str


class Memory(MemoryBase):
    id: UUID
    user_id: UUID
    app_id: UUID
    created_at: datetime
    updated_at: Optional[datetime] = None
    state: str
    categories: Optional[List[Category]] = None
    app: App

    class Config:
        from_attributes = True

class MemoryUpdate(BaseModel):
    content: Optional[str] = None
    metadata_: Optional[dict] = None
    state: Optional[str] = None


class MemoryResponse(BaseModel):
    id: UUID
    content: str
    created_at: int
    state: str
    app_id: UUID
    app_name: str
    categories: List[str]
    metadata_: Optional[dict] = None

    @root_validator(pre=True)
    def process_memory_object(cls, values):
        # If we receive a Memory object, extract the needed fields
        if hasattr(values, '__dict__'):
            # This is a Memory object
            memory_obj = values
            return {
                'id': memory_obj.id,
                'content': memory_obj.content,
                'created_at': int(memory_obj.created_at.timestamp()) if hasattr(memory_obj.created_at, 'timestamp') else memory_obj.created_at,
                'state': memory_obj.state.value if hasattr(memory_obj.state, 'value') else str(memory_obj.state),
                'app_id': memory_obj.app_id,
                'app_name': memory_obj.app.name if hasattr(memory_obj, 'app') and memory_obj.app and hasattr(memory_obj.app, 'name') else "",
                'categories': [cat.name if hasattr(cat, 'name') else str(cat) for cat in memory_obj.categories] if hasattr(memory_obj, 'categories') and memory_obj.categories else [],
                'metadata_': memory_obj.metadata_ if hasattr(memory_obj, 'metadata_') else None
            }
        return values

    @validator('created_at', pre=True)
    def convert_to_epoch(cls, v):
        if isinstance(v, datetime):
            return int(v.timestamp())
        return v
    
    @validator('categories', pre=True)
    def convert_categories(cls, v):
        if v is None:
            return []
        if isinstance(v, list):
            # Handle Category objects
            return [cat.name if hasattr(cat, 'name') else str(cat) for cat in v]
        return v
    
    class Config:
        from_attributes = True

class PaginatedMemoryResponse(BaseModel):
    items: List[MemoryResponse]
    total: int
    page: int
    size: int
    pages: int
